import torch.nn as nn
import torch
from torch.nn.parameter import Parameter

"""
时间：2025
论文名称：Frequency Regulated Channel-Spatial Attention module for improved image classification
论文地址：https://github.com/CoVIS-UNT/FReCSA
"""


class CSHA(nn.Module):
    def __init__(self, inplanes):
        super(CSHA, self).__init__()
        # channel attention
        self.avg_pool_global = nn.AdaptiveAvgPool2d(1)
        self.v = Parameter(torch.zeros(1, inplanes, 1, 1))
        self.batch_norm_channel = nn.BatchNorm2d(inplanes)

        self.sigmoid = nn.Sigmoid()
        self.relu = nn.ReLU(inplace=True)

        # spatial attention
        self.avg_pool = nn.AvgPool2d(kernel_size=7, stride=1, padding=3)
        self.batch_norm_spatial = nn.BatchNorm2d(inplanes)

    def forward(self, x):
        # channel attention
        avg = self.avg_pool_global(x)
        avg = self.batch_norm_channel(avg)
        avg = avg * self.v
        x = x * self.sigmoid(avg)

        low = self.avg_pool(x)
        high = x - low
        s = self.batch_norm_spatial(high * x)
        s = self.relu(s)
        s = self.sigmoid(s)
        x = x * s

        return x


if __name__ == '__main__':
    input = torch.randn(4, 4, 64, 64)
    model = CSHA(4)
    output = model(input)
    print(output.shape)